import numpy as np
import torch
from torch.utils.data import Dataset
import json
import random
from pathlib import Path
import os

class CPCDataset_sameSeq(Dataset):
    def __init__(self, root, n_sample_frames, mode):
        self.root = Path(root)
        self.n_sample_frames = n_sample_frames

        self.speakers = sorted(os.listdir(root/f'{mode}/mels'))
        
        with open(self.root / f"{mode}.json") as file:
            metadata = json.load(file)
        self.metadata = []
        for mel_len, mel_out_path, lf0_out_path,Hubert_path in metadata:
            # if mel_len > n_sample_frames: # only select wavs having frames>=140

            # 实验室集群
            # mel_out_path="/home/wang/codes/py/VC/"+mel_out_path[6:]
            # lf0_out_path="/home/wang/codes/py/VC/"+lf0_out_path[6:]
            # Hubert_path="/home/wang/codes/py/VC/"+Hubert_path[6:]
            
            # AImax2集群
            # mel_out_path="/opt/data/private/VC/"+mel_out_path[6:]
            # lf0_out_path="/opt/data/private/VC/"+lf0_out_path[6:]
            # Hubert_path="/opt/data/private/VC/"+Hubert_path[6:]
            
            # AutoDL
            mel_out_path="/root/autodl-tmp/"+mel_out_path[6:]
            lf0_out_path="/root/autodl-tmp/"+lf0_out_path[6:]
            Hubert_path="/root/autodl-tmp/"+Hubert_path[6:]

            mel_out_path = Path(mel_out_path)
            lf0_out_path = Path(lf0_out_path)
            speaker_path = Hubert_path
            
            # 设置Speaker_emnedding读取路径
            # speaker_path= speaker_path.replace("/root/VCTK/Hubert/"+mode+"/Hubert","/opt/data/private/VC/VCTK/Spkrec-ecapa-voxceleb_Speaker_embedding")
            # speaker_path= speaker_path.replace("/opt/data/private/VC/VCTK/Hubert/"+mode+"/Hubert","/opt/data/private/VC/VCTK/Xvector_Speaker_embedding")
            speaker_path= speaker_path.replace("/root/autodl-tmp/VCTK/Hubert/"+mode+"/Hubert","/root/autodl-tmp/VCTK/Xvector_Speaker_embedding")
            #实验室机器
            # speaker_path= speaker_path.replace("/home/wang/codes/py/VC/VCTK/Hubert/"+mode+"/Hubert","/home/wang/codes/py/VC/VCTK/Xvector_Speaker_embedding")
            speaker_path=Path(speaker_path[:-3]+"flac.npy")
            
            Hubert_path=Path(Hubert_path)
            speaker = mel_out_path.parent.stem
            
            mel_path = self.root.parent / mel_out_path
            lf0_path = self.root.parent / lf0_out_path
            Hubert_path=self.root.parent/Hubert_path
            Speaker_embedding=np.load(speaker_path)
            mel = np.load(mel_path).T
            lf0 = np.load(lf0_path)
            Hubert_vec=np.load(Hubert_path).T
            self.metadata.append([speaker,Speaker_embedding, mel, lf0, Hubert_vec])
        print('n_sample_frames:', n_sample_frames, 'metadata:', len(self.metadata))
        random.shuffle(self.metadata)

    def __len__(self):
        return len(self.metadata)

    def __getitem__(self, index):
        speaker,Speaker_embedding, mel, lf0, Hubert_vec= self.metadata[index]
        print(Speaker_embedding.shape)
        print(Hubert_vec.shape)
       
        melt = mel
        Hubert_vect=Hubert_vec
        lf0t = lf0
        while mel.shape[-1] < self.n_sample_frames:
            mel = np.concatenate([mel, melt], -1)  # 80，T
            lf0 = np.concatenate([lf0, lf0t], 0)
        while Hubert_vec.shape[-1] < self.n_sample_frames//2:
            Hubert_vec=np.concatenate([Hubert_vec,Hubert_vect],-1) #1024,T
            
            
        zero_idxs = np.where(lf0 == 0.0)[0]
        nonzero_idxs = np.where(lf0 != 0.0)[0]
        if len(nonzero_idxs) > 0 :
            mean = np.mean(lf0[nonzero_idxs])
            std = np.std(lf0[nonzero_idxs])
            if std == 0:
                lf0 -= mean
                lf0[zero_idxs] = 0.0
            else:
                lf0 = (lf0 - mean) / (std + 1e-8)
                lf0[zero_idxs] = 0.0
        
        # mel谱归一化
        #注意preprocess_Hubert.py提取特征时，没有对特征进行归一化，这里进行归一化
        # print("mel.shape",mel.shape)
        # print("Hubert.shape",Hubert_vec.shape)
        # mean = np.mean(mel,axis=1).reshape(-1,1)
        # std = np.std(mel,axis=1).reshape(-1,1)
        # mel=(mel - mean) / (std + 1e-8)
        
        # Hubert_vec归一化
        mean = np.mean(Hubert_vec,axis=1).reshape(-1,1)
        std = np.std(Hubert_vec,axis=1).reshape(-1,1)
        Hubert_vec=(Hubert_vec - mean) / (std + 1e-8)

        # pos = random.randint(0, mel.shape[-1] - self.n_sample_frames)
        Hubert_vec_pos = random.randint(0, Hubert_vec.shape[-1] - self.n_sample_frames//2)
        pos = Hubert_vec_pos * 2
        mel = mel[:, pos:pos + self.n_sample_frames]
        lf0 = lf0[pos:pos + self.n_sample_frames]
        # Hubert_vec_pos = random.randint(0, Hubert_vec.shape[-1] - self.n_sample_frames//2)
        # Hubert_vec_pos = pos//2
        Hubert_vec = Hubert_vec[:, Hubert_vec_pos:Hubert_vec_pos + self.n_sample_frames//2]
        # Speaker_embedding
        # print("mel",mel.shape)
        # print("Hubert_vec",Hubert_vec.shape)
        return torch.from_numpy(mel), torch.from_numpy(lf0),torch.from_numpy(Hubert_vec), self.speakers.index(speaker),Speaker_embedding




